import os
import dashscope
from langchain_community.vectorstores import Chroma
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_community.embeddings import DashScopeEmbeddings
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain_core.prompts import ChatPromptTemplate
from data_loader import doc_retriever
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough


OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

client = ChatOpenAI(
    model="gpt-4o-mini",
    api_key=OPENAI_API_KEY
)

# 定义提示模板
message = """
你是问答助手，仅使用检索到的上下文回答下面的问题。
如果不知道答案，就回答不知道
"""

# rag_template = ChatPromptTemplate.from_messages(["human",message])
web_rag_template = ChatPromptTemplate.from_messages([
    ("system",message),
    ("assistant","{context}"),
    ("human","{input}")
])

# 创建文档处理链
combine_docs_chain = create_stuff_documents_chain(
    llm=client,
    prompt=web_rag_template
)

# 创建检索链
retrieval_chain = create_retrieval_chain(
    combine_docs_chain=combine_docs_chain,
    retriever=doc_retriever
)

response = retrieval_chain.invoke({"input": "中美经贸会谈成果是什么？"})
print(response["answer"])